#include "pin.H"
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <fstream>
#include <string>
#include <string.h>
#include <assert.h>
#include <map>
#include <vector>

#include "PinCheckpointer.h"
#include "../HSGlobal.h"

/*TODO: Make the explicit checkpoint capture function take a
        numeric argument that is its index in the checkpoint table.
       
        Create the implicit checkpoint capture function.  Recovery?
*/

PinCheckpointer::PinCheckpointer(vector< void(*)(Recovery *) > *v){

  this->currentCheckPoint = NULL;
  this->activeCheckpoint = 0;
  this->resumingFromCheckPoint = false;
  this->namedCheckpoints.clear();
  this->recoveryRecords.clear();
  this->recoveryRoutines.clear();
  this->maxStack = 0;

  unsigned i; 
  for(i = 0; i < v->size(); i++){

    this->recoveryRoutines[i] = (*v)[i]; 

  }

}


void PinCheckpointer::setActiveCheckpoint(unsigned name){

  this->activeCheckpoint = name;

}

unsigned PinCheckpointer::getActiveCheckpoint(){

  return this->activeCheckpoint;

}

Checkpoint *PinCheckpointer::getCP( unsigned name ){

  if( this->namedCheckpoints.find( name ) == this->namedCheckpoints.end() ){

    this->namedCheckpoints[ name ] = new Checkpoint();

  } 
  return this->namedCheckpoints[ name ];

}


PinCheckpointer *ckptr;

VOID CaptureCheckPoint(CONTEXT *ctx, ADDRINT name, Recovery *r){


  Checkpoint *cp = NULL;
  if(ckptr->resumingFromCheckPoint == true){

    /*The way the instrumentation works out, 
      when we resume from a checkpoint, we
      end up passing through the instrumentation
      that takes a checkpoint. At this point, we've
      set all the registers back to the values
      in the stored context.  That means we can
      go and refill the stack from our stored copy.  
      Of course when resuming it is not useful to
      take another checkpoint, so we just return.
    */
     
    //fprintf(stderr,"In the Capture Check point function with recoveryRoutine=%p\n",r);
    fprintf(stderr,"Getting Checkpoint...\n");
    cp = ckptr->getCP( ckptr->getActiveCheckpoint( ) );

    fprintf(stderr,"Getting Stack Ptr...\n");
    ADDRINT sp = PIN_GetContextReg(ctx, REG_STACK_PTR);

    fprintf(stderr,"copying stack of size %d from %p\n",cp->stackSize,cp->stack);
    memcpy((void*)sp, (void*)cp->stack, cp->stackSize);

    ckptr->resumingFromCheckPoint = false;

    if( r != NULL ){
 
      fprintf(stderr,"In recovery\n");
      fprintf(stderr,"Name: %lu\n",r->name);
      fprintf(stderr,"NumVars: %u\n",r->numVars);
      fprintf(stderr,"vars: {");
      for(unsigned i = 0; i < r->numVars; i++){
        fprintf(stderr,"%p  ",r->vars[i]);
      }
      fprintf(stderr,"}\n"); 

      //Get the recoveryRoutine named in the Recovery from the app
      (ckptr->recoveryRoutines[ r->name ])(r);
      
    }

    return;

  }

  if( r != NULL ){

    unsigned ind = r->name;
    if( ckptr->recoveryRecords.find( ind ) == ckptr->recoveryRecords.end() ){

      /*Copy the recovery routine into Pin's memory space to prevent
        problems if the memory is moved or destroyed due to checkpointing*/
      ckptr->recoveryRecords[ ind ] = (Recovery *)malloc(sizeof(Recovery));
      memcpy( ckptr->recoveryRecords[ind], r, sizeof(Recovery) );

    }

  } 

  cp = ckptr->getCP( name );

  /*Get the stack pointer, copy the contents of the 
    stack to cpStack, our checkpoint buffer*/

  ADDRINT sp = PIN_GetContextReg(ctx, REG_STACK_PTR);

  cp->stackSize = ckptr->maxStack - sp;
  
  if( cp->stack != NULL ){

    free(cp->stack);

  }

  cp->stack = (char *)calloc( cp->stackSize, 1 );

  memcpy((void*)cp->stack, (void*)sp, cp->stackSize);

  PIN_SaveContext( ctx, &(cp->ctx) );

}

VOID RestoreCheckPoint(ADDRINT name){

  /*ExecuteAt restores the register state.
    resumingFromCheckPoint = true makes the 
    checkpoint capturing code (that we will
    jump to after ExecuteAt) restore the stack
    from the "active" checkpoint, set by 
    setActiveCheckpoint*/

  fprintf(stderr,"Restoring checkpoint...%d\n",name);
  ckptr->resumingFromCheckPoint = true;

  fprintf(stderr,"Setting active...%d\n",name);
  ckptr->setActiveCheckpoint( name );

  fprintf(stderr,"Execute At...%d\n",name);
  PIN_ExecuteAt( &(ckptr->getCP( name )->ctx) );

}

VOID UpdateStackSize(ADDRINT memAddr, UINT32 accSize){

  if( memAddr + accSize > ckptr->maxStack ){
  
    ckptr->maxStack = memAddr + accSize;

  }

}





VOID addCheckpointFunctions(RTN rtn, VOID *v){

          RTN_Open(rtn);
          if(strstr(RTN_Name(rtn).c_str(),"LOAD_EXPLICIT_CHECKPOINT")){
              
            RTN_InsertCall(rtn, IPOINT_BEFORE, (AFUNPTR)RestoreCheckPoint, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);

          }
          if(strstr(RTN_Name(rtn).c_str(),"STORE_EXPLICIT_CHECKPOINT_WITH_RECOVERY")){
            RTN_InsertCall(rtn, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_FUNCARG_CALLSITE_VALUE, 1, 
                           IARG_END);
          }

          if(strstr(RTN_Name(rtn).c_str(),"STORE_EXPLICIT_CHECKPOINT") && !strstr(RTN_Name(rtn).c_str(),"STORE_EXPLICIT_CHECKPOINT_WITH_RECOVERY")){
            RTN_InsertCall(rtn, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_UINT32, 0, 
                           IARG_END);
          }
          if(strstr(RTN_Name(rtn).c_str(),"TASK_BEGIN")){
            RTN_InsertCall(rtn, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }
          if(strstr(RTN_Name(rtn).c_str(),"TRY_BEGIN")){
            RTN_InsertCall(rtn, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }
          if(strstr(RTN_Name(rtn).c_str(),"TRY_END")){
            RTN_InsertCall(rtn, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }
          RTN_Close(rtn);

}









VOID pinCheckpointerInstrumentIns(INS ins, VOID *v)
{

  //for( BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) ){

    //for(INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins) ){
      
      if( INS_IsStackWrite(ins) ){

        INS_InsertCall(ins, IPOINT_BEFORE,(AFUNPTR)UpdateStackSize, 
                       IARG_MEMORYWRITE_EA, IARG_MEMORYWRITE_SIZE, 
                       IARG_END);

      }
/*
      if( INS_IsDirectBranchOrCall(ins) && INS_IsProcedureCall(ins) ){

        ADDRINT add = INS_DirectBranchOrCallTargetAddress(ins);

        RTN rtn = RTN_FindByAddress(add);

        if( RTN_Valid(rtn) ){

          if(strstr(RTN_Name(rtn).c_str(),"LOAD_EXPLICIT_CHECKPOINT")){
              
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)RestoreCheckPoint, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
            HSLOG(stderr,"--------------Instrumenting load\n");
          }
          if(strstr(RTN_Name(rtn).c_str(),"STORE_EXPLICIT_CHECKPOINT_WITH_RECOVERY")){
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_FUNCARG_CALLSITE_VALUE, 1, 
                           IARG_END);
            HSLOG(stderr,"---------------Instrumenting with recover\n");
          }

          if(strstr(RTN_Name(rtn).c_str(),"STORE_EXPLICIT_CHECKPOINT") && !strstr(RTN_Name(rtn).c_str(),"STORE_EXPLICIT_CHECKPOINT_WITH_RECOVERY")){
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_UINT32, 0, 
                           IARG_END);
            HSLOG(stderr,"----------------Instrumenting without recover\n");
          }
          if(strstr(RTN_Name(rtn).c_str(),"TASK_BEGIN")){
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }
          if(strstr(RTN_Name(rtn).c_str(),"TRY_BEGIN")){
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }
          if(strstr(RTN_Name(rtn).c_str(),"TRY_END")){
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }

        }

      }*/
/*
    }

  }*/

}

/*To be called in PIN main*/
void initCheckpointer(vector<void(*)(Recovery *)> *v)
{

  fprintf(stderr,"initCheckpointer\n");
  ckptr = new PinCheckpointer(v);
  RTN_AddInstrumentFunction(addCheckpointFunctions, 0);
  INS_AddInstrumentFunction(pinCheckpointerInstrumentIns, 0);

}
